package lock.zookeeper;

import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

import java.io.Closeable;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import static org.apache.zookeeper.CreateMode.EPHEMERAL_SEQUENTIAL;

/**
 * zk独占公平锁
 */
public class ZkclusiveLock implements Closeable {
    private ZooKeeper zk;

    private String lockBasePath = "/lock";
    private String lockPrefix = "exclusive-";
    private String lockName;

    private StringBuffer lock = new StringBuffer();

    public ZkclusiveLock(ZooKeeper zk, String lockName) {
        this.zk = zk;
        this.lockName = lockName;
    }

    public void lock() {
        try {
            // 需要提前创建parent path
            String nodeFullPath = createLockNode();
            while (true) {
                List<String> children = zk.getChildren(getLockPath(), false);
                if (children.isEmpty()) {
                    throw new RuntimeException("lockPath error");
                }
                Collections.sort(children);
                String nodeName = nodeFullPath.substring((getLockPath()).length() + 1);
                if (children.getFirst().equals(nodeName)) {
                    return;
                }
                int index = children.indexOf(nodeName);
                if (index == -1) {
                    // 可能session过期被删了
                    nodeFullPath = createLockNode();
                    continue;
                }
                if (index == 0) {
                    // 获取锁成功
                    return;
                }

                String preNode = children.get(index - 1);
                String preNodePath = getLockPath() + "/" + preNode;

                CountDownLatch latch = new CountDownLatch(1);

                Stat stat = zk.exists(preNodePath, event -> {
                    if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
                        latch.countDown();
                    }
                });
                if (stat == null) {
                    continue;
                }
                // 阻塞等待
                latch.await();
            }

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private String createLockNode() throws KeeperException, InterruptedException {
        String nodeFullPath = zk.create(getLockPath() + "/" + lockPrefix, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, EPHEMERAL_SEQUENTIAL);
        setLock(nodeFullPath);
        return nodeFullPath;
    }

    private void setLock(String nodeFullPath) {
        lock.setLength(0);
        lock.append(nodeFullPath);
    }

    public void unlock() {
        try {
            String name = getName();
            clearName();
            zk.delete(name, -1);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void clearName() {
        setLock("");
    }

    public String getName() {
        return lock.toString();
    }

    public String getLockPath() {
        return lockBasePath + "/" + lockName;
    }

    @Override
    public void close() {
        try {
            zk.close();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    public static void main(String[] args) throws IOException {
        CountDownLatch waitStart = new CountDownLatch(1);
        ZooKeeper zk = new ZooKeeper("localhost:2181", 10 * 1000, event -> {
            if (event.getState() == Watcher.Event.KeeperState.SyncConnected) {
                waitStart.countDown();
            }
        });
        try {
            waitStart.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        ZkclusiveLock lock = new ZkclusiveLock(zk, "product");
        lock.lock();
        System.out.println("get lock");
        lock.unlock();
        lock.close();
        System.out.println("end");
    }
}
